#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" By Martin Senande-Rivera
    For Towards and atmosphere more favourable to firestorm development in Europe """

import os
import glob, sys
import numpy as np
import xarray as xr
import pandas as pd

path_fwi='../0-DATA/FWI/'
path_ki='../0-DATA/Kindex/'
path_outs='./ERA5/'
path_thres='../1-FWI/ERA5/'

from dask_mpi import initialize
initialize()

from distributed import Client
client = Client()

# Read files
fwis=sorted(glob.glob(path_fwi+'ERA5_*.nc'))
FWI = xr.open_mfdataset(fwis,combine='by_coords')['fwi']
FWI = FWI.sel(time=slice('1980-01-01','2018-12-31')).chunk({'time':-1,'latitude': 15, 'longitude': 15})

# Read files
kis=sorted(glob.glob(path_ki+'Ki_*.nc'))
Ki = xr.open_mfdataset(kis,combine='by_coords')['Ki']
Ki = Ki.sel(time=slice('1980-01-01','2018-12-31')).chunk({'time':-1,'latitude': 15, 'longitude': 15})

FWI_thres = xr.open_dataset(path_thres+'FWI_90threshold.nc')['fwi'] # FWI threshold
Ki_thres = xr.full_like(FWI_thres,25.)                              # Kindex threshold

Ki['time']=FWI['time'] # Skip possible differences in time dimension deffinition

num = xr.full_like(Ki, 0).chunk({'time':-1,'latitude': 15, 'longitude': 15}) # Daily array full of 0
num = xr.where((FWI>FWI_thres) & (Ki>Ki_thres),1,num)                        # Days with FWI and Ki above each threshold are classified as 1
ndays = num.chunk({'time':-1,'latitude': 15, 'longitude': 15}).resample(time='1Y').sum(dim='time') # Number of days with FWI and Ki above each threshold
ndays.chunk({'time':-1,'latitude': 15, 'longitude': 15}).to_netcdf(path_outs+'FWI-Ki_ndays.nc')    # Save output

ndays_p = ndays.sel(time=slice('1980-01-01','2018-12-31')).chunk({'time':-1,'latitude': 15, 'longitude': 15}) # Select the 1980-2018 period
ndays_present = ndays_p.sum(dim='time')/39.                    # Average number of days per year with FWI and Ki above each threshold in the 1980-2018 period
ndays_present.to_netcdf(path_outs+'FWI-Ki_ndays-present.nc')   # Save output
